import numpy as np
import pandas as pd
from sklearn.datasets import load_iris, load_breast_cancer, load_wine
from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import GaussianNB, MultinomialNB, BernoulliNB
from sklearn.metrics import accuracy_score, confusion_matrix
import matplotlib.pyplot as plt

# 加载数据集
iris = load_iris()
breast_cancer = load_breast_cancer()
wine = load_wine()

# 定义朴素贝叶斯模型
models = {
    "Gaussian Naive Bayes": GaussianNB(),
    "Multinomial Naive Bayes": MultinomialNB(),
    "Bernoulli Naive Bayes": BernoulliNB()
}

datasets = {
    "Iris": iris,
    "Breast Cancer": breast_cancer,
    "Wine": wine
}

# 循环遍历数据集和模型
for dataset_name, dataset in datasets.items():
    X = dataset.data
    y = dataset.target

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

    plt.figure(figsize=(15, 5))

    for i, (model_name, model) in enumerate(models.items(), 1):
        clf = model.fit(X_train, y_train)
        y_pred = clf.predict(X_test)

        accuracy = accuracy_score(y_test, y_pred)
        confusion = confusion_matrix(y_test, y_pred)

        plt.subplot(1, len(models), i)
        plt.imshow(confusion, cmap=plt.cm.Blues)
        plt.title(f"{model_name}\nAccuracy: {accuracy:.2f}")
        plt.colorbar()
        tick_marks = np.arange(len(dataset.target_names))
        plt.xticks(tick_marks, dataset.target_names, rotation=45)
        plt.yticks(tick_marks, dataset.target_names)
        plt.xlabel('Predicted Label')
        plt.ylabel('True Label')
        for i in range(len(dataset.target_names)):
            for j in range(len(dataset.target_names)):
                plt.text(j, i, confusion[i, j], ha='center', va='center', color='black')

    plt.suptitle(f"Comparison of Naive Bayes Models on {dataset_name} Dataset")
    plt.tight_layout()
    plt.show()
